#!/usr/bin/python
# -*- coding: utf-8 -*-
from __future__ import print_function
""" crypto.cipher.tkip_encr

    TKIP encryption from IEEE 802.11 TGi

    TKIP uses WEP (ARC4 with crc32) and key mixing
    This is only the encryption and not the Michael integrity check!

    Copyright © (c) 2002 by Paul A. Lambert
    Read LICENSE.txt for license information.

    November 2002
"""
from crypto.cipher.arc4 import ARC4
from zlib import crc32
from struct import pack
from crypto.keyedHash.tkip_key_mixing import TKIP_Mixer
from crypto.errors import BadKeySizeError, IntegrityCheckError
from binascii_plus import *


class TKIP_encr:
    """ TKIP Stream Cipher Algorithm without the Michael integrity check

        TKIP encryption on an MPDU using WEP with a longer 'iv'
        and the TKIP key mixing algorithm .  This does NOT include
        the Michael integrity algorithm that operates on the MSDU data.
    """

    def __init__(self, key=None, transmitterAddress=None, keyID=None):
        """ Initialize TKIP_encr, key -> octet string for key """
        assert(keyID == 0 or keyID == None), 'keyID should be zero in TKIP'
        self.keyId = 0
        self.name = 'TKIP_encr'
        self.strength = 128
        self.encryptHeaderSize = 8        # used to skip octets on decrypt
        self.arc4 = ARC4()                # base algorithm
        self.keyMixer = TKIP_Mixer(key, transmitterAddress)
        if key != None:                   # normally in base init, uses adjusted keySize
            self.setKey(key)
        if transmitterAddress != None:
            self.setTA(transmitterAddress)

    def setKey(self, key, ta=None):
        """ Set key, key string is 16 octets long """
        self.keyMixer.setKey(key)
        if ta != None:
            self.setTA(ta)

    def setTA(self, transmitterAddress):
        """ Set the transmitter address """
        self.keyMixer.setTA(transmitterAddress)

    def _getIVandKeyID(self, cipherText):
        """ Parse the TKIP header to get iv and set KeyID
            iv is returned as octet string and is little-endian!!!
        """
        assert(ord(cipherText[3]) & 0x20), 'extIV SHOULD be set in TKIP header'
        self.setCurrentKeyID = (ord(cipherText[3]) & 0xC0) >> 6
        return cipherText[:3] + cipherText[5:9] # note iv octets are little-endian!!!

    def _makeARC4key(self, tscOctets, keyID=0):
        """ Make an ARC4 key from TKIP Sequence Counter Octets (little-endian) """
        if keyID != 0:
            raise 'TKIP expects keyID of zero'
        print("tscOctets in tkmixer=", b2a_p(tscOctets))
        newKey = self.keyMixer.newKey(tscOctets)
        print("newKey=", b2a_p(newKey))
        return newKey

    def encrypt(self, plainText, iv):
        """ Encrypt a string and return a binary string
            iv is 6 octets of little-endian encoded pn
        """
        assert(len(iv) == 6), 'TKIP bad IV size on encryption'
        self.pnField = iv
        self.arc4.setKey(self._makeARC4key(iv))
        eh1 = chr((ord(iv[0]) | 0x20) & 0x7f)
        encryptionHeader = iv[0] + eh1 + iv[1] + chr((self.keyId << 6) | 0x20) + iv[2:]
        crc = pack('<I', crc32(plainText))
        cipherText = encryptionHeader + self.arc4.encrypt(plainText + crc)
        return cipherText

    def decrypt(self, cipherText):
        """ Decrypt a WEP packet, assumes WEP 4 byte header on packet """
        assert(ord(cipherText[3]) & 0x20), 'extIV SHOULD be set in TKIP header'
        self.setCurrentKeyID = (ord(cipherText[3]) & 0xC0) >> 6
        iv = cipherText[0] + cipherText[2] + cipherText[4:8]
        self.pnField = iv
        self.arc4.setKey(self._makeARC4key(iv))
        plainText = self.arc4.decrypt(cipherText[self.encryptHeaderSize:])
        if plainText[-4:] != pack('<I', crc32(plainText[:-4])):  # check data integrity
            raise IntegrityCheckError('WEP CRC Integrity Check Error')
        return plainText[:-4]
